#include<stdio.h>
#include<stdlib.h>
#include<math.h>
#include<string.h>
#include<mpi.h>

void verify(double *y_cpu, double *y, int size) {
    double eps = 1e-9;
    int error_cnt = 0;
    for (int i = 0; i < size; i++) {
        if (fabs(y_cpu[i] - y[i]) > eps) {
            error_cnt++;
        }
    }
    if(error_cnt > 0) {
        printf("Error! Total %d errors found!\n", error_cnt);
    } else {
        printf("Verification passed!\n");
    }
}

void MPI_Allreduce_ring(void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype, MPI_Op op, MPI_Comm comm) {
    int rank, nproc;
    MPI_Comm_rank(comm, &rank);
    MPI_Comm_size(comm, &nproc);
    double *sbuf = (double*)sendbuf;
    double *rbuf = (double*)recvbuf;
    int chunk_size = count / nproc;
    int remainder = count % nproc;
    MPI_Request send_req, recv_req;
    double *temp_buf = (double *)malloc(chunk_size * sizeof(double));

    // ScatterReduce
    for(int step = 0; step < nproc - 1; step++) {
        int send_idx = (rank - step + nproc) % nproc;
        int send_buf_idx = send_idx * chunk_size;
        if(send_idx == nproc - 1) {
            chunk_size += remainder;    
        }
        MPI_Isend(sbuf + send_buf_idx, chunk_size, datatype, (rank + 1) % nproc, 0, comm, &send_req);

        int recv_idx = (rank - step - 1 + nproc) % nproc;
        int recv_buf_idx = recv_idx * chunk_size;
        if(recv_idx == nproc - 1) {
            chunk_size += remainder;
        }
        MPI_Irecv(temp_buf, chunk_size, datatype, (rank - 1 + nproc) % nproc, 0, comm, &recv_req);

        MPI_Wait(&send_req, MPI_STATUS_IGNORE);
        MPI_Wait(&recv_req, MPI_STATUS_IGNORE);

        for(int i = 0; i < chunk_size; i++) {
            if(op == MPI_SUM) {
                sbuf[recv_buf_idx + i] += temp_buf[i];
            } else if(op == MPI_MAX) {
                if(temp_buf[i] > sbuf[recv_buf_idx + i]) {
                    sbuf[recv_buf_idx + i] = temp_buf[i];
                }
            }
        }
    }

    // Allgather
    for(int step = 0; step < nproc - 1; step++) {
        int send_idx = (rank - step + 1 + nproc) % nproc;
        int send_buf_idx = send_idx * chunk_size;
        if(send_idx == nproc - 1) {
            chunk_size += remainder;    
        }
        MPI_Isend(sbuf + send_buf_idx, chunk_size, datatype, (rank + 1) % nproc, 0, comm, &send_req);

        int recv_idx = (rank - step + nproc) % nproc;
        int recv_buf_idx = recv_idx * chunk_size;
        if(recv_idx == nproc - 1) {
            chunk_size += remainder;
        }
        MPI_Irecv(sbuf + recv_buf_idx, chunk_size, datatype, (rank - 1 + nproc) % nproc, 0, comm, &recv_req);

        MPI_Wait(&send_req, MPI_STATUS_IGNORE);
        MPI_Wait(&recv_req, MPI_STATUS_IGNORE);
    }
    memcpy(rbuf, sbuf, count * sizeof(double));
    free(temp_buf);
}

int main(int argc, char **argv) {
    int rank, nproc;
    int m;
    double *x, *y_mpi, *y_ring;
    MPI_Op op;
    char opt_str[256];

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &nproc);

    if(rank == 0) {
        FILE *fp1;
        fp1 = fopen("input2.txt", "r");
        fscanf(fp1, "%d,%s", &m, opt_str);
        fclose(fp1);
        
        if(strcmp(opt_str, "SUM") == 0) {
            op = MPI_SUM;
        } else if(strcmp(opt_str, "MAX") == 0) {
            op = MPI_MAX;
        }
    }

    MPI_Bcast(&m, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(opt_str, 256, MPI_CHAR, 0, MPI_COMM_WORLD);

    if(rank != 0) {
        if(strcmp(opt_str, "SUM") == 0) {
            op = MPI_SUM;
        } else if(strcmp(opt_str, "MAX") == 0) {
            op = MPI_MAX;
        }
    }  

    x = (double *)malloc(m * sizeof(double));
    y_mpi = (double *)malloc(m * sizeof(double));
    y_ring = (double *)malloc(m * sizeof(double));
    for(int i = 0; i < m; i++) {
        x[i] = (double)rand() / RAND_MAX;
    }

    double mpi_start = MPI_Wtime();
    MPI_Allreduce(x, y_mpi, m, MPI_DOUBLE, op, MPI_COMM_WORLD);
    double mpi_end = MPI_Wtime();
    double mpi_time = mpi_end - mpi_start;

    MPI_Barrier(MPI_COMM_WORLD);
    double ring_start = MPI_Wtime();
    MPI_Allreduce_ring(x, y_ring, m, MPI_DOUBLE, op, MPI_COMM_WORLD);
    double ring_end = MPI_Wtime();
    double ring_time = ring_end - ring_start;

    if(rank == 0) {
        verify(y_mpi, y_ring, m); 
        printf("MPI_Allreduce time: %f seconds\n", mpi_time);
        printf("Ring Allreduce time: %f seconds\n", ring_time);
        printf("speedup: %f\n", mpi_time / ring_time);

        FILE *fp2 = fopen("output2.txt", "w");
        fprintf(fp2, "%.2f", mpi_time);
        fprintf(fp2, "%.2f", ring_time);
        fclose(fp2);
    }
    free(x);
    free(y_mpi);
    free(y_ring);
    MPI_Finalize();
    return 0;
}